import PickBlue
import ImageEnhance
import LineCut
import cv2
import calScaleRed
import calScaleBlack


def removeRed(origin_img, img, canny_img):
    res_img = img.copy()
    h, w = img.shape[:2]
    for i in range(0,h):
        for j in range(0,w):
            B = float(origin_img[i][j][0])
            G = float(origin_img[i][j][1])
            R = float(origin_img[i][j][2])
            # if canny_img[i][j] > 0:
            #     # findNum = False
            #     # for m in range(max(i - 1, 0), min(i + 1, h - 1)):
            #     #     for n in range(max(0, j - 1), min(j + 1, w - 1)):
            #     #         if img[m][n] == 255:
            #     #             findNum = True
            #     #             break
            #     # if findNum:
            #     res_img[i][j] = 255
            if (R / max(G + R + B, 1)) * 255 > 93 and R+B+G > 315:
                res_img[i][j] = 0

    # for i in range(0, h):
    #     for j in range(0, w):
    #         res_img[i][j] = 255 - res_img[i][j]
    # res_img = PickBlue.erode_demo(res_img, 2)
    # # res_img = PickBlue.dilate_demo(res_img, 2)
    # for i in range(0, h):
    #     for j in range(0, w):
    #         res_img[i][j] = 255 - res_img[i][j]
    return res_img


if __name__ == "__main__":
    img_path = 'receipt_img/ticket1.jpg'
    img = cv2.imread(img_path)
    h, w = img.shape[:2]
    print("H&W: ", h, w)
    scaleRed = calScaleRed.calScale(img)
    scaleBlack = calScaleBlack.calScale(img)
    print("计算相对大小，红黑尺寸：", scaleRed, scaleBlack)
    resScale = h/30
    if scaleRed < h/5:
        resScale = scaleRed/4.6
        print("选择红色印章作为相对尺寸大小")
    elif scaleBlack < h/15:
        resScale = scaleBlack
        print("选择黑色标志作为相对尺寸大小")
    else:
        print("红黑标志计算相对尺寸不适用，使用默认值")
    print("相对大小结果", int(resScale + 0.5), "\n")
    print("----------开始进行图像处理----------")
    cut_images = PickBlue.pickBlue(img, int(resScale + 0.5))
    label_1, label_2 = 0, 0
    for i in cut_images:
        print("\n------------------", end="")
        print("收据" + str(label_1) + "------------------")
        # i = ImageEnhance.imageEnhance(i)
        line_imgs = LineCut.wordCut(i)
        label_2 = 0
        for j in line_imgs:
            enhanceImg = ImageEnhance.imageEnhance(j)
            line_canny = cv2.Canny(enhanceImg, 0, 100)
            # if label_1==0 and label_2<5:
            #     cv2.imshow("line_canny"+str(label_1)+"_"+str(label_2), line_canny)
            #     cv2.imshow("line_img" + str(label_1) + "_" + str(label_2), enhanceImg)
            ret, thresholdImg = cv2.threshold(enhanceImg, 150, 255, cv2.THRESH_BINARY_INV)
            thresholdImg = removeRed(j, thresholdImg, line_canny)
            # thresholdImg = PickBlue.dilate_demo(thresholdImg, 2)
            # thresholdImg = cv2.GaussianBlur(thresholdImg, (3,3),0)
            # if label_1 == 0 and label_2 == 5:
            #     cv2.imshow("ori_line", j)
            #     cv2.imwrite("cutLines/" + "ori_" + str(label_1) + "_" + str(label_2) + ".jpg", j)
            #     cv2.imshow("en_line", enhanceImg)
            #     cv2.imwrite("cutLines/" + "en_" + str(label_1) + "_" + str(label_2) + ".jpg", enhanceImg)
            #     cv2.imshow("img_"+str(label_1)+"_"+str(label_2), thresholdImg)
            #     cv2.imwrite("cutLines/"+"img_"+str(label_1)+"_"+str(label_2)+".jpg",thresholdImg)
            cv2.imshow("img_" + str(label_1) + "_" + str(label_2), thresholdImg)
            cv2.imwrite("cutLines/" + "img_" + str(label_1) + "_" + str(label_2) + ".jpg", thresholdImg)
            label_2 += 1
        label_1 += 1

    print("############ 完成 finished ############")
    cv2.waitKey(0)
    cv2.destroyAllWindows()
